#!/usr/bin/python3

""" 
    Skeleton code for k-means clustering mini-project.
"""

import os
import joblib
import numpy as np
import matplotlib.pyplot as plt
import sys
sys.path.append(os.path.abspath("../tools/"))
from feature_format import featureFormat, targetFeatureSplit
from sklearn.cluster import KMeans
from sklearn.preprocessing import MinMaxScaler

def Draw(pred, features, poi, mark_poi=False, name="image.png", f1_name="feature 1", f2_name="feature 2"):
    """ some plotting code designed to help you visualize your clusters """

    ### plot each cluster with a different color--add more colors for
    ### drawing more than five clusters
    colors = ["b", "r", "k", "m", "g"]
    for ii, pp in enumerate(pred):
        plt.scatter(features[ii][0], features[ii][1], color = colors[pred[ii]])

    ### if you like, place red stars over points that are POIs (just for funsies)
    if mark_poi:
        for ii, pp in enumerate(pred):
            if poi[ii]:
                plt.scatter(features[ii][0], features[ii][1], color="r", marker="*")
    plt.xlabel(f1_name)
    plt.ylabel(f2_name)
    plt.savefig(name)
    plt.show()

def scale_features(d, enable=False):
    if not enable:
        return d
    scaler = MinMaxScaler()
    scaler.fit(d)
    # print(scaler.transform( [[ 200000., 1000000.]] ))
    return scaler.transform(d)

def show_precluster(plotter, fin_data, no_of_features):
    if no_of_features == 2:
        for f1, f2 in fin_data:
            plotter.scatter( f1, f2 )
        plotter.show()
    elif no_of_features == 3:
        for f1, f2,_ in fin_data:
            plotter.scatter( f1, f2 )
        plotter.show()

def main(data_dict, more_features=[], feature_scale=False):

    ### the input features we want to use 
    ### can be any key in the person-level dictionary (salary, director_fees, etc.) 
    feature_1 = "salary"
    feature_2 = "exercised_stock_options"
    features_list = ["poi", feature_1, feature_2]
    features_list += more_features
    data = featureFormat(data_dict, features_list )
    poi, finance_features = targetFeatureSplit( data )

    ### in the "clustering with 3 features" part of the mini-project,
    ### you'll want to change this line to 
    ### for f1, f2, _ in finance_features:
    ### (as it's currently written, the line below assumes 2 features)
    n = len(features_list)- 1

    show_precluster(plt, finance_features, n)

    ### cluster here; 
    ### create predictions of the cluster labels for the data
    ### and store them to a list called pred
    kmeans = KMeans(
        n_clusters=n
    )
    kmeans.fit(scale_features(finance_features, enable=feature_scale))
    pred = kmeans.labels_ 

    try:
        Draw(pred, finance_features, poi, 
             mark_poi=False, 
             name="%iclusters.pdf" % n, 
             f1_name=feature_1, 
             f2_name=feature_2)
    except NameError:
        print("No predictions object named pred found, no clusters to plot")

def get_range_for(dataset, key):
    values  = [ dataset[person][key] for person in dataset \
        if type(dataset[person][key]) == int ]
    # print(values)
    # ordered = values.sort(reverse=True)
    # if ordered == None:
    #     print("got none from sorting")
    # else:
    print("(min, max) = (%i, %i)" % (max(values), min(values)))
    

if __name__ == "__main__":
        ### load in the dict of dicts containing all the data on each person in the dataset
    DATASET_PATH = "../final_project/final_project_dataset.pkl"
    dd = joblib.load( open(DATASET_PATH, "rb") )
    ### there's an outlier--remove it! 
    dd.pop("TOTAL", 0)

    m = ["total_payments"]
    main(dd, more_features=[], feature_scale=True)

    get_range_for(dd, "exercised_stock_options")
    get_range_for(dd, "salary")

